"""
A2A协议服务器实现 - 支持标准A2A方法
"""
import json
import logging
from typing import Any, Dict, Optional, AsyncGenerator
from fastapi import APIRouter,FastAPI, Request
from fastapi.responses import JSONResponse, StreamingResponse, FileResponse
import os

from .types import (
    AgentCard, JSONRPCRequest, JSONRPCResponse, JSONRPCError,
    SendTaskRequest, GetTaskRequest, TaskNotFoundError,
    InvalidRequestError, InternalError, MethodNotFoundError,
    Message, MessageSendParams, Task, TaskStatusUpdateEvent, TaskArtifactUpdateEvent, TaskStatus, TaskState
)
from .task_manager import TaskManager, AITaskManager
from uuid import uuid4
from datetime import datetime

logger = logging.getLogger(__name__)

api_router = APIRouter()

class DateTimeEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, datetime):
            return obj.isoformat()
        return super().default(obj)

# 辅助类：异步生成器资源管理器
class AsyncGeneratorResourceManager:
    def __init__(self):
        self.generators = []

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        # 反向关闭所有注册的生成器
        for gen in reversed(self.generators):
            try:
                await gen.aclose()
                pass
            except Exception as e:
                logger.error(f"Error closing generator: {e}")
        return False  # 不抑制异常

    def register(self, generator):
        self.generators.append(generator)
        return generator

class A2AServer:
    """A2A协议服务器"""
    
    def __init__(
        self,
        agent_card: AgentCard,
        task_manager: TaskManager,
        host: str = "0.0.0.0",
        port: int = None,
        lifespan = None
    ):
        self.agent_card = agent_card
        self.task_manager = task_manager
        self.host = host
        self.port = port
        self.app = FastAPI(
            title=agent_card.name,
            description=agent_card.description,
            version=agent_card.version,
            lifespan=lifespan
        )

        self.app.include_router(api_router)

        # 注册路由
        self._setup_routes()

    def _setup_routes(self):
        """设置路由"""
        
        @self.app.get("/.well-known/agent.json")
        async def get_agent_card():
            """获取智能体卡片"""
            return self.agent_card.model_dump(exclude_none=True)
        
        @self.app.post("/api/v1/a2a/")
        async def handle_a2a_request(request: Request):
            """处理A2A JSON-RPC请求"""
            try:
                body = await request.json()
                logger.info(f"Received A2A request: {body}")
                return await self._process_request(body)
                    
            except json.JSONDecodeError:
                return self._create_error_response(
                    None, InvalidRequestError(message="Invalid JSON")
                )
            except Exception as e:
                logger.error(f"Error processing request: {e}")
                return self._create_error_response(
                    None, InternalError(message=str(e))
                )

        # 新增图片接口，返回图片流数据
        @self.app.get("/api/v1/icon")
        async def get_image(name: str):
            """获取图片流数据"""
            try:
                # 这里需要实现根据 image_id 找到图片文件路径的逻辑
                # 示例假设图片存放在 ./images 目录下
                image_dir = f"./{name}/images"
                image_path = os.path.join(image_dir, f"icon.png")

                if os.path.exists(image_path):
                    return FileResponse(image_path)
                else:
                    return self._create_error_response(
                        None, InternalError(message="Image not found")
                    )
            except Exception as e:
                logger.error(f"Error getting image: {e}")
                return self._create_error_response(
                    None, InternalError(message=str(e))
                )

        @self.app.get("/api/v1/file")
        async def get_file(path: str,name: str):
            """获取图片流数据"""
            try:
                # 这里需要实现根据 name 找到文件路径的逻辑
                # 示例假设文件存放在 ./docs 目录下
                file_dir = f"./{path}/docs"
                file_path = os.path.join(file_dir, f"{name}")

                if os.path.exists(file_path):
                    return FileResponse(file_path)
                else:
                    return self._create_error_response(
                        None, InternalError(message="File not found")
                    )
            except Exception as e:
                logger.error(f"Error getting file: {e}")
                return self._create_error_response(
                    None, InternalError(message=str(e))
                )

    async def _process_request(self, body: Dict[str, Any]):
        """统一处理JSON-RPC请求"""
        try:
            # 基本验证
            if not isinstance(body, dict):
                return self._create_error_response(
                    None, InvalidRequestError(message="Request must be an object")
                )
            
            if body.get("jsonrpc") != "2.0":
                return self._create_error_response(
                    body.get("id"), InvalidRequestError(message="Invalid JSON-RPC version")
                )
            
            method = body.get("method", "")
            params = body.get("params", {})
            request_id = body.get("id")
            
            logger.info(f"Processing method: {method}")
            
            # 根据方法类型决定返回格式
            if method == "message/stream":
                # 流式方法：返回StreamingResponse
                return await self._create_streaming_response(request_id, params)
            else:
                # 非流式方法：返回JSONResponse
                return await self._create_json_response(request_id, method, params)
            
        except Exception as e:
            logger.error(f"Error processing request: {e}")
            return self._create_error_response(
                body.get("id") if isinstance(body, dict) else None,
                InternalError(message=str(e))
            )
    
    async def _create_streaming_response(self, request_id: str, params: Dict[str, Any]) -> StreamingResponse:
        """创建流式响应"""
        async def stream_generator():
            try:
                # 发送SSE连接建立的初始消息
                yield "event: connected\ndata: {\"status\": \"connected\"}\n\n"
                async_gen = self._handle_message_stream(params)
                async for event in async_gen:
                    # 创建JSON-RPC流式响应
                    response = JSONRPCResponse(id=request_id, result=event)
                    response_json = response.model_dump(exclude_none=True)
                    
                    # 使用标准SSE格式，确保立即发送
                    event_data = f"event: status-update\ndata: {json.dumps(response_json, ensure_ascii=False, cls=DateTimeEncoder)}\n\n"
                    yield event_data
                    
                    # 如果是最终事件，结束流
                    if not isinstance(event, TaskArtifactUpdateEvent) and event.final:
                        # 发送流结束标记
                        yield "event: end\ndata: {\"status\": \"completed\"}\n\n"
                        break
                #if async_gen:
                #    await async_gen.aclose()
            except Exception as e:
                logger.error(f"Error in streaming processing: {e}")
                # 发送错误事件
                error_response = JSONRPCResponse(
                    id=request_id, 
                    error=InternalError(message=str(e))
                )
                error_json = error_response.model_dump(exclude_none=True)
                yield f"event: error\ndata: {json.dumps(error_json, ensure_ascii=False)}\n\n"
        
        return StreamingResponse(
            stream_generator(),
            media_type="text/event-stream",  # 使用正确的SSE媒体类型
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "Content-Type": "text/event-stream; charset=utf-8",
                "Access-Control-Allow-Origin": "*",  # 允许跨域
                "Access-Control-Allow-Headers": "Cache-Control",
                "X-Accel-Buffering": "no",  # 禁用Nginx缓冲
            }
        )
    
    async def _create_json_response(self, request_id: str, method: str, params: Dict[str, Any]) -> JSONResponse:
        """创建JSON响应"""
        # 路由到相应的处理方法
        if method == "message/send":
            # 标准A2A消息发送方法
            result = await self._handle_message_send(params)
        elif method == "tasks/send":
            # 兼容旧的方法
            result = await self.task_manager.handle_send_task(params)
        elif method == "tasks/get":
            result = await self.task_manager.handle_get_task(params)
        else:
            return self._create_error_response(
                request_id, MethodNotFoundError(message=f"Method '{method}' not found")
            )
        
        # 创建成功响应
        response = JSONRPCResponse(id=request_id, result=result)
        return JSONResponse(response.model_dump(exclude_none=True))
    
    async def _handle_message_send(self, params: Dict[str, Any]) -> Task:
        """处理标准A2A message/send方法"""
        try:
            # 解析MessageSendParams
            logger.info(f"Processing message/send request: {params}")
            
            # 提取消息信息
            message_data = params.get("message", {})
            configuration = params.get("configuration", {})
            
            # 生成sessionId作为contextId使用
            session_id = message_data.get('sessionId','') or str(uuid4())
            
            # 转换为内部任务格式
            task_params = {
                "id": params.get("id", str(uuid4())),
                "sessionId": session_id,  # 确保有sessionId用于contextId生成
                "message": message_data,
                "metadata": params.get("metadata")
            }
            
            # 使用任务管理器处理
            task = await self.task_manager.handle_send_task(task_params)
            
            logger.info(f"message/send processing completed: {task}")
            return task
            
        except Exception as e:
            logger.error(f"Error processing message/send: {e}")
            raise
    
    async def _handle_message_stream(self, params: Dict[str, Any]) -> AsyncGenerator[Any, None]:
        """处理标准A2A message/stream方法 - 真正的流式返回"""
        try:
            logger.info(f"Processing message/stream request: {params}")
            
            # 检查任务管理器是否支持流式处理
            if not hasattr(self.task_manager, 'execute_task_streaming'):
                raise ValueError("任务管理器不支持流式处理，缺少execute_task_streaming方法")
            
            # 提取消息信息
            message_data = params.get("message", {})
            configuration = params.get("configuration", {})
            
            # 转换为内部任务格式
            task_params = {
                "id": params.get("id", ""),
                "message": message_data,
                "configuration": configuration
            }

            # 创建任务
            from .types import TaskSendParams, Message
            send_params = TaskSendParams(**task_params)
            task = await self.task_manager.create_task(send_params)

            #async with AsyncGeneratorResourceManager() as resource_manager:
            # 流式执行任务
            event_stream = self.task_manager.execute_task_streaming(task, send_params.message)
            #resource_manager.register(event_stream)
            async for event in event_stream:
                if isinstance(event, TaskStatusUpdateEvent):
                    # 提取消息文本内容用于日志
                    message_text = ""
                    if event.status.message and event.status.message.parts:
                        from .task_manager import extract_text_from_part
                        message_text = extract_text_from_part(event.status.message.parts[0])
                        # 限制日志中显示的文本长度
                        if len(message_text) > 100:
                            message_text = message_text[:100] + "..."

                    logger.info(f"Streaming event: kind={event.kind}, state={event.status.state}, final={event.final}, text='{message_text}'")
                elif isinstance(event, TaskArtifactUpdateEvent):
                    # 处理TaskArtifactUpdateEvent类型的事件
                    logger.info(f"Artifact update: taskId={event.taskId}, artifacts={event.artifact}")
                elif isinstance(event, Message):
                    # 处理Message类型的事件
                    message_text = ""
                    if event.parts:
                        from .task_manager import extract_text_from_part
                        message_text = extract_text_from_part(event.parts[0])
                        # 限制日志中显示的文本长度
                        if len(message_text) > 100:
                            message_text = message_text[:100] + "..."
                    logger.info(f"Streaming message: text='{message_text}'")
                yield event

        except Exception as e:
            logger.error(f"Error processing message/stream: {e}")
            raise
    
    def _create_error_response(self, request_id: Any, error: JSONRPCError) -> JSONResponse:
        """创建错误响应"""
        response = JSONRPCResponse(id=request_id, error=error)
        return JSONResponse(
            response.model_dump(exclude_none=True),
            status_code=400
        )
    
    def get_app(self) -> FastAPI:
        """获取FastAPI应用实例"""
        return self.app